﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using NetPE.Core.Pe;
using NetPE.TreeNodes;
using System.Windows.Forms;
using System.Drawing;
using NetPE.Core;
using System.Reflection;
using System.Globalization;
using NetPE.Core.DataDirectories;

namespace NetPE
{
    public class DirectoryLocationEditor : BaseEditor
    {
        DirectoryNode n;
        TableLayoutPanel panel;
        public DirectoryLocationEditor(DirectoryNode n)
        {
            this.n = n;
            this.DoubleBuffered = true;
            panel = new TableLayoutPanel();
            panel.ColumnCount = 2;
            panel.Location = Point.Empty;
            panel.AutoSize = true;
            panel.ColumnStyles.Add(new ColumnStyle(SizeType.Absolute, 200));
            panel.ColumnStyles.Add(new ColumnStyle(SizeType.Absolute, 200));
            panel.CellBorderStyle = TableLayoutPanelCellBorderStyle.Single;
            panel.Padding = new Padding(5);
            this.Controls.Add(panel);
        }

        protected override void RefreshData()
        {
            panel.Controls.Clear();

            PropertyInfo[] info = new PropertyInfo[] { typeof(VirtualLocation).GetProperty("Address"), typeof(VirtualLocation).GetProperty("Size") };
            panel.RowCount = 3;

            Button load = new Button() { Text = "Load", Dock = DockStyle.Fill, FlatStyle = FlatStyle.Popup };
            load.Click += new EventHandler(Load);
            Button save = new Button() { Text = "Save", Dock = DockStyle.Fill, FlatStyle = FlatStyle.Popup };
            save.Click += new EventHandler(Save);
            panel.Controls.Add(load, 0, 0);
            panel.Controls.Add(save, 1, 0);
            
            for (int i = 0; i < info.Length; i++)
            {
                panel.Controls.Add(new Label() { Text = info[i].Name, Dock = DockStyle.Fill, AutoSize = false, TextAlign = ContentAlignment.MiddleCenter }, 0, i + 1);

                TextBox tb;
                switch (info[i].PropertyType.FullName)
                {
                    case "NetPE.Core.Rva":
                        tb = new TextBox();
                        tb.BackColor = Color.Black;
                        tb.ForeColor = Color.White;
                        tb.BorderStyle = BorderStyle.None;
                        tb.TextAlign = HorizontalAlignment.Center;
                        tb.Width = 200;
                        tb.Anchor = AnchorStyles.None;
                        tb.Text = info[i].GetValue(Object, null).ToString();
                        tb.LostFocus += new EventHandler(PropChanged);
                        tb.Tag = info[i];
                        panel.Controls.Add(tb, 1, i + 1);
                        break;
                    case "System.UInt32":
                        tb = new TextBox();
                        tb.BackColor = Color.Black;
                        tb.ForeColor = Color.White;
                        tb.BorderStyle = BorderStyle.None;
                        tb.TextAlign = HorizontalAlignment.Center;
                        tb.Width = 200;
                        tb.Anchor = AnchorStyles.None;
                        tb.Text = ((uint)info[i].GetValue(Object, null)).ToString("x8");
                        tb.LostFocus += new EventHandler(PropChanged);
                        tb.Tag = info[i];
                        panel.Controls.Add(tb, 1, i + 1);
                        break;
                }
            }

            panel.RowStyles.Add(new RowStyle(SizeType.Absolute, 32));
            panel.RowStyles.Add(new RowStyle(SizeType.Absolute, 32));
            panel.RowStyles.Add(new RowStyle(SizeType.Absolute, 32));
        }

        bool onevent = false;
        private void PropChanged(object sender, EventArgs e)
        {
            if (!onevent)
            {
                PropertyInfo prop = (sender as Control).Tag as PropertyInfo;
                string val = (sender as Control).Text;

                switch (prop.PropertyType.FullName)
                {
                    case "NetPE.Core.Rva":
                        if (val.TrimStart('0').Length > 8) break;
                        prop.SetValue(Object, (Rva)uint.Parse(val, NumberStyles.HexNumber), null);
                        break;
                    case "System.UInt32":
                        if (val.TrimStart('0').Length > 8) break;
                        prop.SetValue(Object, uint.Parse(val, NumberStyles.HexNumber), null);
                        break;
                    default: throw new InvalidOperationException();
                }

                onevent = true;
                switch (prop.PropertyType.FullName)
                {
                    case "NetPE.Core.Rva":
                        (sender as Control).Text = prop.GetValue(Object, null).ToString();
                        break;
                    case "System.UInt32":
                        (sender as Control).Text = ((uint)prop.GetValue(Object, null)).ToString("x8");
                        break;
                    default: throw new InvalidOperationException();
                }
                onevent = false;
            }
        }

        void Save(object sender, EventArgs e)
        {
            if (n.Nodes.Count == 0) return;
            DataDirectory dd = (n.Nodes[0] as BaseNode).Data as DataDirectory;
            if (dd.Type == DataDirectoryType.IAT ||
                dd.Type == DataDirectoryType.Certificate)
                return;

            PeFile f = (n.Parent.Parent as FileNode).Data as PeFile;

            uint size = dd.GetTotalSize();
            if ((n.Data as DataDirectoryEntry).TotalSize != size && MessageBox.Show("Old total Size : " + (n.Data as DataDirectoryEntry).TotalSize + "\nNew total Size : " + size + "\nDifferent Size! Continue?", "NetPE", MessageBoxButtons.YesNo, MessageBoxIcon.Question) == DialogResult.No)
                return;

            size = dd.GetDirectorySize();
            if ((n.Data as DataDirectoryEntry).Size != size && MessageBox.Show("Old directory size : " + (n.Data as DataDirectoryEntry).Size + "\nNew directory Size : " + size + "\nDifferent Size! Continue?", "NetPE", MessageBoxButtons.YesNo, MessageBoxIcon.Question) == DialogResult.No)
                return;

            
            dd.Save(new VirtualWriter(f.SectionHeaders.GetVirtualStream()));
        }

        void Load(object sender, EventArgs e)
        {
            if (n.Nodes.Count != 0 && MessageBox.Show("Loaded Directory! Continue?", "NetPE", MessageBoxButtons.YesNo, MessageBoxIcon.Question) == DialogResult.No)
                return;
            n.Nodes.Clear();
            
            PeFile f = (n.Parent.Parent as FileNode).Data as PeFile;

            DataDirectory dd;
            if ((n.Data as DataDirectoryEntry).Type == DataDirectoryType.IAT)
            {
                MessageBox.Show("Use Import Directory instead of IAT Directory!!");
                return;
            }
            else if ((n.Data as DataDirectoryEntry).Type == DataDirectoryType.Certificate)
            {
                dd = f.Certificates;
            }
            else
            {
                try
                {
                    (n.Data as DataDirectoryEntry).Reload();
                    dd = (n.Data as DataDirectoryEntry).GetComponent() as DataDirectory;
                    dd.Load(new VirtualReader(f.SectionHeaders.GetVirtualStream()));
                }
                catch(NotSupportedException)
                {
                    MessageBox.Show("Not support directory!");
                    return;
                }
                catch (InvalidOperationException)
                {
                    MessageBox.Show("Invaild directory!");
                    return;
                }
            }

            BaseNode bn;
            switch (dd.Type)
            {
                case DataDirectoryType.Certificate:
                    bn = new CertificateDNode(dd); break;
                case DataDirectoryType.CLR:
                    bn = new CLRDNode(dd); break;
                case DataDirectoryType.Debug:
                    bn = new DebugDNode(dd); break;
                case DataDirectoryType.Export:
                    bn = new ExportDNode(dd); break;
                case DataDirectoryType.Import:
                    bn = new ImportDNode(dd); break;
                case DataDirectoryType.Relocation:
                    bn = new RelocDNode(dd); break;
                case DataDirectoryType.Resource:
                    bn = new ResourceDNode(dd); break;
                case DataDirectoryType.TLS:
                    bn = new TLSDNode(dd); break;
                default:
                    MessageBox.Show("Not exist directory editor!");
                    return;
            }
            n.Nodes.Add(bn);
        }
    }
}
